/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.apache.spark.network.netty

import java.nio.ByteBuffer

import scala.collection.JavaConverters._
import scala.language.existentials
import scala.reflect.ClassTag

import org.apache.spark.internal.Logging
import org.apache.spark.network.BlockDataManager
import org.apache.spark.network.buffer.{ManagedBuffer, NioManagedBuffer}
import org.apache.spark.network.client.{RpcResponseCallback, TransportClient}
import org.apache.spark.network.server.{OneForOneStreamManager, RpcHandler, StreamManager}
import org.apache.spark.network.shuffle.protocol.{BlockTransferMessage, OpenBlocks, StreamHandle, UploadBlock}
import org.apache.spark.serializer.Serializer
import org.apache.spark.storage.{BlockId, StorageLevel}

/**
 * Serves requests to open blocks by simply registering one chunk per block requested.
 * Handles opening and uploading arbitrary BlockManager blocks.
 *
 * Opened blocks are registered with the "one-for-one" strategy, meaning each Transport-layer Chunk
 * is equivalent to one Spark-level shuffle block.
 *
  * @param appId 当前应用的Application ID
  * @param serializer 序列化器
  * @param blockManager 所属的BlockManager
  */
class NettyBlockRpcServer(
    appId: String,
    serializer: Serializer,
    blockManager: BlockDataManager)
  extends RpcHandler with Logging {

  // 提供一对一的流服务
  private val streamManager = new OneForOneStreamManager()

  override def receive(
      client: TransportClient,
      rpcMessage: ByteBuffer,
      responseContext: RpcResponseCallback): Unit = {
    // 获取并解码消息
    val message = BlockTransferMessage.Decoder.fromByteBuffer(rpcMessage)
    logTrace(s"Received request: $message")

    // 根据消息类型分别处理
    message match {
      case openBlocks: OpenBlocks => // 打开Block的消息
        // 取出消息携带的BlockId数组，获取每个BlockId对应的Block的ManagedBuffer，封装为序列
        val blocks: Seq[ManagedBuffer] =
          openBlocks.blockIds.map(BlockId.apply).map(blockManager.getBlockData)
        // 使用OneForOneStreamManager对象将这些Block的ManagedBuffer序列注册到streams缓存
        val streamId = streamManager.registerStream(appId, blocks.iterator.asJava)
        logTrace(s"Registered streamId $streamId with ${blocks.size} buffers")
        // 返回StreamHandle消息，包含Stream ID和ManagedBuffer序列的大小
        responseContext.onSuccess(new StreamHandle(streamId, blocks.size).toByteBuffer)

      case uploadBlock: UploadBlock => // 上传Block的消息
        // StorageLevel and ClassTag are serialized as bytes using our JavaSerializer.
        // 对消息携带的元数据进行反序列化，得到存储级别和类型标记
        val (level: StorageLevel, classTag: ClassTag[_]) = {
          serializer
            .newInstance()
            .deserialize(ByteBuffer.wrap(uploadBlock.metadata))
            .asInstanceOf[(StorageLevel, ClassTag[_])]
        }
        // 将UploadBlock消息携带的Block数据（即blockData），封装为NioManagedBuffer。
        val data = new NioManagedBuffer(ByteBuffer.wrap(uploadBlock.blockData))
        // 获取UploadBlock消息携带的BlockId。
        val blockId = BlockId(uploadBlock.blockId)
        // 调用BlockManager的putBlockData方法，将Block存入本地存储体系。
        blockManager.putBlockData(blockId, data, level, classTag)
        // 通过响应上下文回复客户端。
        responseContext.onSuccess(ByteBuffer.allocate(0))
    }
  }

  // 返回OneForOneStreamManager对象
  override def
  getStreamManager(): StreamManager = streamManager
}
